package cn.xxi.ai.workflow.core.graph;

import cn.xxi.ai.workflow.core.node.ConditionNode;
import cn.xxi.ai.workflow.core.edge.Edge;
import cn.xxi.ai.workflow.core.node.NormalNode;
import cn.xxi.ai.workflow.core.node.Node;
import cn.xxi.ai.workflow.core.node.NodeData;
import cn.xxi.ai.workflow.core.node.end.EndNode;
import cn.xxi.ai.workflow.core.node.start.StartNode;
import cn.xxi.ai.workflow.state.State;
import org.bsc.langgraph4j.CompiledGraph;
import org.bsc.langgraph4j.GraphStateException;
import org.bsc.langgraph4j.StateGraph;
import org.bsc.langgraph4j.action.AsyncEdgeAction;
import org.bsc.langgraph4j.action.AsyncNodeAction;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 表示一个完整的工作流图，包括节点和连边。
 */
public class Graph {

    private final List<Node<? extends NodeData>> nodes;
    private final List<Edge> edges;

    public Graph(List<Node<? extends NodeData>> nodes, List<Edge> edges) {
        this.nodes = nodes;
        this.edges = edges;
    }

    /**
     * 校验并返回合法的起始节点（StartNode）。
     *
     * @throws IllegalAccessException 如果没有起始节点，或有多个起始节点，或者起始节点没有连出边
     */
    public StartNode validStartNode() throws IllegalAccessException {
        List<StartNode> startNodes = nodes.stream()
                .filter(n -> n instanceof StartNode)
                .map(n -> (StartNode) n)
                .toList();

        if (startNodes.isEmpty()) {
            String message = "missing start node.";
            throw new IllegalAccessException(message);
        }

        if (startNodes.size() > 1) {
            String message = "only one start node allowed, but found " + startNodes.size() + ".";
            throw new IllegalAccessException(message);
        }

        StartNode startNode = startNodes.getFirst();
        boolean hasOutgoing = edges.stream().anyMatch(e -> e.getSource().equals(startNode.getId()));
        if (!hasOutgoing) {
            String message = "start node must connect to at least one node.";
            throw new IllegalAccessException(message);
        }

        return startNode;
    }

    /**
     * 校验并返回合法的结束节点（EndNode）列表。
     *
     * @throws IllegalAccessException 如果没有结束节点，或者所有结束节点都没有被连入
     */
    public List<EndNode> validEndNode() throws IllegalAccessException {
        List<EndNode> endNodes = nodes.stream()
                .filter(n -> n instanceof EndNode)
                .map(n -> (EndNode) n)
                .collect(Collectors.toList());

        if (endNodes.isEmpty()) {
            String message = "missing end node.";
            throw new IllegalAccessException(message);
        }

        boolean hasConnectedEnd = endNodes.stream()
                .anyMatch(n -> edges.stream().anyMatch(e -> e.getTarget().equals(n.getId())));

        if (!hasConnectedEnd) {
            String message = "end node must connect to at least one node.";
            throw new IllegalAccessException(message);
        }

        return endNodes;
    }

    /**
     * 将当前 Graph 编译成 LangGraph4j 的可执行图（CompiledGraph<State>）。
     */
    public CompiledGraph<State> compile() throws IllegalAccessException, GraphStateException {
        StartNode startNode = validStartNode();
        List<EndNode> endNodes = validEndNode();

        List<NormalNode<?>> normalNodes = nodes.stream()
                .filter(n -> n instanceof NormalNode)
                .map(n -> (NormalNode<?>) n)
                .collect(Collectors.toList());

        List<ConditionNode<?>> conditionNodes = nodes.stream()
                .filter(n -> n instanceof ConditionNode)
                .map(n -> (ConditionNode<?>) n)
                .collect(Collectors.toList());

        // 构建普通边（排除条件节点的连接）
        Set<String> conditionNodeIds = conditionNodes.stream()
                .map(Node::getId)
                .collect(Collectors.toSet());

        List<Edge> normalEdges = edges.stream()
                .filter(e -> !conditionNodeIds.contains(e.getSource()))
                .filter(e -> !conditionNodeIds.contains(e.getTarget()))
                .collect(Collectors.toList());

        // 添加起始节点和结束节点的连接边
        normalEdges.add(new Edge(StateGraph.START, startNode.getId()));
        for (EndNode endNode : endNodes) {
            normalEdges.add(new Edge(endNode.getId(), StateGraph.END));
        }

        StateGraph<State> graph = new StateGraph<>(State::new);

        // 添加普通节点
        for (NormalNode<?> node : normalNodes) {
            graph.addNode(node.getId(), AsyncNodeAction.node_async(node));
        }

        // 添加普通边
        for (Edge edge : normalEdges) {
            graph.addEdge(edge.getSource(), edge.getTarget());
        }

        // 特殊处理条件节点
        for (ConditionNode<?> node : conditionNodes) {
            List<Edge> targetEdges = edges.stream()
                    .filter(e -> e.getSource().equals(node.getId()))
                    .toList();

            List<Edge> sourceEdges = edges.stream()
                    .filter(e -> e.getTarget().equals(node.getId()))
                    .toList();

            for (Edge edge : sourceEdges) {
                Map<String, String> branches = targetEdges.stream()
                        .collect(Collectors.toMap(Edge::getSourceHandle, Edge::getTarget));

                graph.addConditionalEdges(
                        edge.getSource(),
                        AsyncEdgeAction.edge_async(node),
                        branches
                );
            }
        }

        return graph.compile();
    }
}
